# -*- coding: utf-8 -*-
"""
Created on Wed Sep 11 15:35:29 2013

@author: BigSk_000


Paper: High-resolution correlation spectroscopy of 13C spins near a nitrogen-vacancy centre in diamond
"""

from qutip import *
from pylab import *
import nv_hamilton
import qutip_enhanced
import numpy as np
import coordinates as co

qte = qutip_enhanced.QutipEnhanced()

#initial state
rho0 = tensor(fock_dm(2,0), fock_dm(2,0))

#hamilton
magic_angle = 54.74/360.0*2*pi
nvh = nv_hamilton.NVHam(magnet_field = {'z':0.55}, n_type = None, electron_levels = [1,2])
nvh.add_spin(nvh.hf_tensor_13c_dd(location = {'rho': 0.5e-9, 'elev': np.pi/2}), nvh.h_13c())

rho = qte.rotate(rho0, rotation_axis = {'x': 1}, rotated_spin = 0, angle = np.pi/2)

#print e[4] - e[5]
#print e[6] - e[7]


#phi = pi/2
#rho0 = tensor(1/2.0*(fock_dm(2,0) + fock_dm(2,1)))
#o = qu_tools.Operator(op = rho0)
#a = Bloch3d()
#a.add_states(rho0)
#a.show()
#rho = o.rotate(rotation_axis = {'azim': phi}, angle = pi/2, rotated_spin = 0, transition = [0,1])
#b = Bloch3d()
#b.add_states(rho)
#b.show()
#o = qu_tools.Operator(op = rho)
#rho = o.rotate(rotation_axis = {'y': 1}, angle = pi/2, rotated_spin = 0, transition = [0,1])
#c = Bloch3d()
#c.add_states(rho)
#c.show()


##NV hamilton operator
#spins = [1,1/2.0]
#
##hyperfine 

#c13_location = {'rho':6e-10, 'elev':magic_angle} #corresponds to ~92kHz hyperfine splitting
#h_hf = nv_hamilton.NVHam().get_h_hf(nv_hamilton.NVHam().hf_tensor_13c_dd(c13_location), spins)
#hf_splitting = abs(h_hf.diag()[0] - h_hf.diag()[1])
#print hf_splitting
##larmor
#gamma_13c = 10.705
#factor_larmor = 1
#larmor = factor_larmor * hf_splitting
#B_z = larmor/gamma_13c
#tau = 1/(2*gamma_13c*B_z)*1e6 #1/2*larmor, *1e6 because frequencies in megahertz, so timescales in microseconds
#H = B_z*gamma_13c * tensor(qeye(3), jmat(1/2.0, 'z')) + h_hf
#
##start density matrix
#rho0 = tensor(fock_dm(3, 0), fock_dm(2, 0)) #initialize electron and n15 in lowest energy state
#
#rho = qu_tools.rotate_spin(rho0, axis = 'y', spins = spins, rotated_spin = 0, angle = pi/2.0)
#rho = mesolve(H, rho, np.linspace(0,tau,20), [], []).states[-1]

#rho0 = tensor(fock_dm(3, 0), fock_dm(2, 0)) #initialize electron and n15 in lowest energy state
#rho = qu_tools.rotate_spin(rho0, rotation_axis = {'y':1, 'x':2}, spins = [1,1/2.0], rotated_spin = 0, angle = pi, transition = [0,1])
#b = Bloch3d()
#b.add_states(rho.ptrace(1))
#b.show()


#b = Bloch3d()
#b.add_states(o.states[-1])
#b.show()

#sequence
#rho = qu_tools.rotate_spin(rho0, axis = 'y', spins = [1,1/2.0], rotated_spin = 1, angle = pi/4)



#b = Bloch3d()
#b.add_states(rho.ptrace(1))
#rho = qu_tools.rotate_spin(rho0, axis = 'y', spins = [1,1/2.0], rotated_spin = 1, angle = pi/4)
#b.show()
#    delta = 0.0 * 2 * pi  # qubit sigma_x coefficient
#    eps0 = 1.0 * 2 * pi   # qubit sigma_z coefficient
#    A = 0.25 * 2 * pi     # drive amplitude (reducing -> RWA more accurate)
#    w = 1.0 * 2 * pi      # drive frequency
#    gamma1 = 0.0          # relaxation rate
#    n_th = 0.0            # average number of excitations ("temperature")
#    psi0 = basis(2, 1)    # initial state
#    
#    #
#    # Hamiltonian
#    #
#    sx = sigmax()
#    sz = sigmaz()
#    sm = destroy(2)
#    
#    H0 = - delta / 2.0 * sx - eps0 / 2.0 * sz
#    H1 = - A * sx
#    
#    print H0    
#    
#    # define the time-dependence of the hamiltonian using the list-string
#    # format
#    args = {'w': w}
#    Ht = [H0, [H1, 'sin(w*t)']]
#    
#    print Ht    
#    
#    #
#    # collapse operators
#    #
#    c_op_list = []
#    
#    rate = gamma1 * (1 + n_th)
#    if rate > 0.0:
#        c_op_list.append(sqrt(rate) * sm)       # relaxation
#    
#    rate = gamma1 * n_th
#    if rate > 0.0:
#        c_op_list.append(sqrt(rate) * sm.dag())  # excitation
#    
#    #
#    # evolve and system subject to the time-dependent hamiltonian
#    #
#    tlist = linspace(0, 5.0 * 2 * pi / A, 500)
#    output1 = mesolve(Ht, psi0, tlist, c_op_list, [sm.dag() * sm], args)
#    
#    # Alternative: write the hamiltonian in a rotating frame, and neglect the
#    # the high frequency component (rotating wave approximation), so that the
#    # resulting Hamiltonian is time-independent.
#    H_rwa = - delta / 2.0 * sx - A * sx / 2
#    output2 = mesolve(H_rwa, psi0, tlist, c_op_list, [sm.dag() * sm])
#    
#    #
#    # Plot the solution
#    #
#    plot(tlist, real(output1.expect[0]), 'b',
#         tlist, real(output2.expect[0]), 'r')
#    xlabel('Time')
#    ylabel('Occupation probability')
#    title('Excitation probabilty of qubit')
#    legend(("Time-dependent Hamiltonian", "Corresponding RWA"))
#    show()
#    close()